# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import json

from absl import (flags, app)
from coin.strategy.basis_strat.sim.sim_job_generator import run_on_slurm

FLAGS = flags.FLAGS


def _generate_config(param_config, keys, i, config):
  if i < len(keys):
    key = keys[i]
    values = param_config[key]
    for val in values:
      config[key] = val
      if i == len(keys) - 1:
        yield config.copy()
      else:
        for config in _generate_config(param_config, keys, i + 1, config):
          yield config.copy()


def generate_all_combination(param_config):
  keys = sorted(list(param_config.keys()))
  for config in _generate_config(param_config, keys, 0, {}):
    yield config


def enumerate_dates(start_date, end_date):
  cur_date = start_date
  while cur_date <= end_date:
    yield cur_date
    cur_date += datetime.timedelta(days=1)


def get_configs(config_path):
  with open(config_path, 'r') as fi:
    config = json.load(fi)
    return config['param_config'], config['default_config']


def main(argv):
  assert FLAGS.trading_date
  assert FLAGS.config_file

  trading_dates = []
  if '-' in FLAGS.trading_date:
    start_date, end_date = FLAGS.trading_date.split('-')
    trading_dates = [
        d for d in enumerate_dates(datetime.datetime.strptime(start_date, '%Y%m%d'),
                                   datetime.datetime.strptime(end_date, '%Y%m%d'))
    ]
    trading_dates = [int(d.strftime('%Y%m%d')) for d in trading_dates]
  else:
    trading_dates = [int(FLAGS.trading_date)]

  param_config, default_config = get_configs(FLAGS.config_file)
  assert 'exchanges' in default_config
  assert 'base' in default_config
  assert 'quote' in default_config
  assert 'sub_dir' in default_config
  param_config['trading_date'] = trading_dates
  default_config['output_dir'] = os.path.join(FLAGS.output_dir, default_config['sub_dir'])
  del default_config['sub_dir']
  configs = []
  for config in generate_all_combination(param_config):
    config.update(default_config)
    configs.append(config)
  run_on_slurm('spread_basis', 'coin.experimental.yuxuan.spread_basis_strategy', configs, mem=512)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', '', 'YYYYMMDD or YYYYMMDD-YYYYMMDD')
  flags.DEFINE_string('output_dir', '/remote/iosg/home/yuxuan/coin/sim_result', '')
  app.run(main)
